package shell

import (
	"bytes"
	"github.com/fatih/color"
	"golang.org/x/crypto/ssh"
	"gosh/config"
	"gosh/pkg/connect"
	"gosh/pkg/handler"
	"strings"
)

var jumpServer *connect.JumpServer

func Run() {
	if err := initJumpServer(); err != nil {
		color.Red("init bastion failed,err:%s\n", err.Error())
		return
	}
	if jumpServer != nil {
		defer jumpServer.Close()
	}
	handler.Handler(config.SSHInfo.Fork, config.SSHInfo.Hosts, config.SSHInfo.PrintStatus, config.SSHInfo.PrintFields, process)
}

func initJumpServer() error {
	if config.SSHInfo.BastionHost != "" {
		jumpServer = connect.NewJumpServer(config.SSHInfo.BastionHost, config.SSHInfo.BastionUser,
			config.SSHInfo.BastionPassword, config.SSHInfo.BastionKey, config.SSHInfo.BastionKeyPassphrase,
			config.SSHInfo.BastionPort, config.SSHInfo.ConnectTimeout)
		return jumpServer.Open()
	}
	return nil
}

func process(ip string) (stdout, stderr []byte, err error) {
	host := connect.NewHost(ip, config.SSHInfo.User, config.SSHInfo.Password, config.SSHInfo.Key,
		config.SSHInfo.KeyPassphrase, config.SSHInfo.Port, jumpServer, config.SSHInfo.ConnectTimeout)
	if err = host.Open(); err != nil {
		return nil, nil, err
	}
	defer host.Close()
	cmd := config.SSHInfo.Args
	return execCommand(host.SSHClient, strings.Join(cmd, " "))
}

func execCommand(client *ssh.Client, cmd string) (stdout, stderr []byte, err error) {
	session, err := client.NewSession() // 开启新的ssh会话
	if err != nil {
		return nil, nil, err
	}
	defer func() {
		_ = session.Close()
		_ = client.Close()
	}()
	// 指定标准输出和标准错误
	var stdOut bytes.Buffer
	var stdErr bytes.Buffer
	session.Stderr = &stdErr
	session.Stdout = &stdOut
	err = session.Run(cmd)
	return stdOut.Bytes(), stdErr.Bytes(), err
}
